"""
STORK (Stabilized Runge-Kutta) sampler implementation for ComfyUI
Based on the paper: "STORK: Improving the Fidelity of Mid-NFE Sampling for Diffusion and Flow Matching Models"
"""

import math
import torch
import numpy as np
from scipy.io import loadmat
from pathlib import Path
import comfy.samplers
from comfy.k_diffusion import sampling as k_diffusion_sampling

# Path to STORK parameter files
STORK_PARAMS_PATH = Path(__file__).parent / "stork_parameters"

class STORKSampler:
    """STORK sampler for both diffusion and flow matching models"""
    
    def __init__(self, solver_order=2, s=50, derivative_order=2, prediction_type="epsilon"):
        """
        Initialize STORK sampler
        
        Args:
            solver_order: 2 or 4 for STORK2 or STORK4
            s: Number of sub-steps
            derivative_order: Order of Taylor expansion (2 or 3)
            prediction_type: "epsilon" for diffusion or "flow_prediction" for flow matching
        """
        self.solver_order = solver_order
        self.s = s
        self.derivative_order = derivative_order
        self.prediction_type = prediction_type
        
        # Storage for derivative approximations
        self.model_outputs = []
        self.velocity_predictions = []
        self.noise_predictions = []
        
        # Load STORK4 coefficients if needed
        if solver_order == 4:
            self._load_stork4_coefficients()
            
    def _load_stork4_coefficients(self):
        """Load precomputed STORK4 coefficients from .mat files"""
        try:
            ms_path = STORK_PARAMS_PATH / "ms.mat"
            fpa_path = STORK_PARAMS_PATH / "fpa.mat"
            fpb_path = STORK_PARAMS_PATH / "fpb.mat"
            recf_path = STORK_PARAMS_PATH / "recf.mat"
            
            self.ms = loadmat(str(ms_path))['ms']
            self.fpa = loadmat(str(fpa_path))['fpa']
            self.fpb = loadmat(str(fpb_path))['fpb']
            self.recf = loadmat(str(recf_path))['recf'].flatten()
        except Exception as e:
            print(f"Warning: Could not load STORK4 coefficients: {e}")
            # Fall back to basic coefficients
            self.ms = None
            self.fpa = None
            self.fpb = None
            self.recf = None
            
    def b_coeff(self, j):
        """Calculate b_j coefficients for STORK2"""
        if j == 0:
            return 1
        elif j == 1:
            return 1 / 2
        else:
            return ((2 * j - 3) / (2 * j)) * self.b_coeff(j - 1)
            
    def mdegr(self, s, ms):
        """Find the degree and parameters for STORK4"""
        if ms is None:
            # Fallback if coefficients not loaded
            return min(s, 50), [0, 0]
            
        mdeg = 0
        for i in range(ms.shape[0]):
            if s >= ms[i, 0]:
                mdeg = int(ms[i, 0])
                mp = ms[i, 1:3]
                break
        return mdeg, mp
        
    def taylor_approximation(self, order, diff, v0, v1, v2=None, v3=None):
        """Taylor expansion approximation for velocity/noise"""
        if order == 1:
            return v0
        elif order == 2:
            return v0 + diff * v1 + (diff ** 2) / 2 * v2
        elif order == 3:
            return v0 + diff * v1 + (diff ** 2) / 2 * v2 + (diff ** 3) / 6 * v3
        else:
            return v0
            
    def startup_phase(self, model_output, sample, sigma, sigma_next):
        """Startup phase for derivative approximation"""
        self.model_outputs.append(model_output)
        
        if self.prediction_type == "flow_prediction":
            self.velocity_predictions.append(model_output)
            # Simple Euler step for startup
            return sample - (sigma - sigma_next) * model_output
        else:  # epsilon prediction
            self.noise_predictions.append(model_output)
            # Convert epsilon to velocity for diffusion
            sigma_dims = sigma.ndim
            sigma_expanded = sigma
            while len(sigma_expanded.shape) < len(model_output.shape):
                sigma_expanded = sigma_expanded.unsqueeze(-1)
            velocity = -model_output / sigma_expanded
            return sample + (sigma_next - sigma) * velocity


def append_dims(x, ndim):
    """Append dimensions to match target dimensionality"""
    if x.ndim == ndim:
        return x
    return x.view(*x.shape, *([1] * (ndim - x.ndim)))


def sample_stork2(model, x, sigmas, extra_args=None, callback=None, disable=None, s=50, derivative_order=2, debug=False):
    """
    STORK2 sampler for ComfyUI
    
    Args:
        model: The model function
        x: Initial noise/sample
        sigmas: Noise schedule
        extra_args: Additional arguments for the model
        callback: Callback function
        disable: Disable progress bar
        s: Number of sub-steps
        derivative_order: Order of Taylor expansion
        debug: Enable debug output
    """
    extra_args = {} if extra_args is None else extra_args
    
    if debug:
        print(f"\n=== STORK2 Sampler Debug ===")
        print(f"Input shape: {x.shape}")
        print(f"Input dtype: {x.dtype}")
        print(f"Input device: {x.device}")
        print(f"Sigmas shape: {sigmas.shape}")
        print(f"Sigmas range: [{sigmas.min():.6f}, {sigmas.max():.6f}]")
        print(f"Actual sigma values (first 5): {sigmas[:5].tolist()}")
        print(f"Number of steps: {len(sigmas) - 1}")
        print(f"Substeps (s): {s}")
        print(f"Derivative order: {derivative_order}")
        print(f"Extra args keys: {extra_args.keys() if extra_args else 'None'}")
    
    sampler = STORKSampler(solver_order=2, s=s, derivative_order=derivative_order)
    
    # Determine if this is flow matching or diffusion
    is_flow_matching = hasattr(model, 'inner_model') and hasattr(model.inner_model, 'prediction_type')
    if is_flow_matching:
        sampler.prediction_type = "flow_prediction"
    
    if debug:
        print(f"Model type detection: {'Flow Matching' if is_flow_matching else 'Diffusion (Epsilon)'}")
        print(f"Prediction type: {sampler.prediction_type}")
        if hasattr(model, 'inner_model'):
            print(f"Inner model type: {type(model.inner_model).__name__}")
    
    s_in = x.new_ones([x.shape[0]])
    
    for i in range(len(sigmas) - 1):
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        
        if sigma_next == 0:
            # Skip the last step if sigma_next is 0
            if debug:
                print(f"Step {i}: Skipping final step (sigma_next=0)")
            break
            
        # Get model output
        denoised = model(x, sigma * s_in, **extra_args)
        
        if debug and i < 3:  # Only debug first few steps
            print(f"\nStep {i}:")
            print(f"  sigma: {sigma:.6f}, sigma_next: {sigma_next:.6f}")
            print(f"  x stats - mean: {x.mean():.4f}, std: {x.std():.4f}, min: {x.min():.4f}, max: {x.max():.4f}")
            print(f"  denoised stats - mean: {denoised.mean():.4f}, std: {denoised.std():.4f}, min: {denoised.min():.4f}, max: {denoised.max():.4f}")
        
        # Convert denoised to derivative (velocity)
        # In k-diffusion, the model returns denoised, and we convert to ODE derivative
        d = (x - denoised) / append_dims(sigma, x.ndim)
        
        if debug and i < 3:
            print(f"  ODE derivative stats: mean: {d.mean():.4f}, std: {d.std():.4f}")
            sigma_expanded = append_dims(sigma, x.ndim)
            print(f"  Sigma expanded shape: {sigma_expanded.shape}, values: {sigma_expanded.flatten()[0]:.6f}")
        
        # Startup phase for first few steps
        if i < sampler.derivative_order:
            sampler.model_outputs.append(d)
            # Simple Euler step for startup
            dt = sigmas[i + 1] - sigmas[i]
            x = x + d * dt
            if debug and i < 3:
                print(f"  Startup phase - Euler step with dt={dt:.6f}")
        else:
            # Main STORK2 step
            dt = sigma_next - sigma  # dt should be negative (going from high to low sigma)
            
            # Calculate derivatives
            h1 = sigmas[i-1] - sigmas[i]
            h2 = sigmas[i-2] - sigmas[i-1] if i >= 2 else h1
            
            outputs = sampler.model_outputs
            
            if sampler.derivative_order == 2 and len(outputs) >= 2:
                # 2nd order derivative approximation
                derivative = (-outputs[-2] + 4 * outputs[-1] - 3 * d) / (2 * h1)
                second_derivative = 2 / (h1 * h2 * (h1 + h2)) * (outputs[-2] * h1 - outputs[-1] * (h1 + h2) + d * h2)
                third_derivative = None
            else:
                # Fallback to 1st order
                derivative = (d - outputs[-1]) / h1 if len(outputs) >= 1 else torch.zeros_like(d)
                second_derivative = torch.zeros_like(d)
                third_derivative = None
            
            sampler.model_outputs.append(d)
            
            # STORK2 sub-stepping
            Y_j_2 = x
            Y_j_1 = x
            Y_j = x
            
            if debug and i < 3:
                print(f"  Starting STORK2 substeps with dt={dt:.6f}")
                print(f"  h1={h1:.6f}, h2={h2:.6f}")
                print(f"  derivative stats: mean={derivative.mean():.4f}, std={derivative.std():.4f}")
                if second_derivative is not None:
                    print(f"  second_derivative stats: mean={second_derivative.mean():.4f}, std={second_derivative.std():.4f}")
            
            for j in range(1, sampler.s + 1):
                if j == 1:
                    mu_tilde = 6 / ((sampler.s + 4) * (sampler.s - 1))
                    Y_j = Y_j_1 + dt * mu_tilde * d
                    if debug and i < 3 and j == 1:
                        print(f"    Substep {j}: mu_tilde={mu_tilde:.6f}, dt*mu_tilde*d mean={( dt * mu_tilde * d).mean():.6f}")
                else:
                    fraction = ((j - 1)**2 + (j - 1) - 2) / (sampler.s**2 + sampler.s - 2) if j > 2 else 4 / (3 * (sampler.s**2 + sampler.s - 2))
                    
                    mu = (2 * j + 1) * sampler.b_coeff(j) / (j * sampler.b_coeff(j - 1))
                    nu = -(j + 1) * sampler.b_coeff(j) / (j * sampler.b_coeff(j - 2))
                    mu_tilde = mu * 6 / ((sampler.s + 4) * (sampler.s - 1))
                    gamma_tilde = -mu_tilde * (1 - j * (j + 1) * sampler.b_coeff(j-1) / 2)
                    
                    diff = fraction * dt
                    velocity = sampler.taylor_approximation(sampler.derivative_order, diff, d, derivative, second_derivative, third_derivative)
                    
                    if debug and i < 3 and j in [2, sampler.s]:  # Debug substep 2 and last substep
                        print(f"    Substep {j}: mu={mu:.4f}, nu={nu:.4f}, mu_tilde={mu_tilde:.6f}, gamma_tilde={gamma_tilde:.6f}")
                        print(f"      fraction={fraction:.6f}, diff={diff:.6f}")
                        print(f"      velocity mean={velocity.mean():.6f}, Y_j_1 mean={Y_j_1.mean():.6f}")
                    
                    Y_j = mu * Y_j_1 + nu * Y_j_2 + (1 - mu - nu) * x + dt * mu_tilde * velocity + dt * gamma_tilde * d
                
                Y_j_2 = Y_j_1
                Y_j_1 = Y_j
            
            x = Y_j
            
            if debug and i < 3:
                print(f"  After STORK2 substeps - x stats: mean: {x.mean():.4f}, std: {x.std():.4f}")
            
        # Keep only recent outputs for memory efficiency
        if len(sampler.model_outputs) > sampler.derivative_order + 1:
            sampler.model_outputs.pop(0)
            
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigma, 'sigma_hat': sigma, 'denoised': denoised})
        
        # Debug the progression every 5 steps
        if debug and i % 5 == 0:
            print(f"\nProgress - Step {i}/{len(sigmas)-1}:")
            print(f"  Current x stats: mean={x.mean():.4f}, std={x.std():.4f}, min={x.min():.4f}, max={x.max():.4f}")
            
    if debug:
        print(f"\nFinal output:")
        print(f"  x stats: mean={x.mean():.4f}, std={x.std():.4f}, min={x.min():.4f}, max={x.max():.4f}")
            
    return x


def sample_stork4(model, x, sigmas, extra_args=None, callback=None, disable=None, s=50, derivative_order=2, debug=False):
    """
    STORK4 sampler for ComfyUI
    
    Args:
        model: The model function
        x: Initial noise/sample
        sigmas: Noise schedule
        extra_args: Additional arguments for the model
        callback: Callback function
        disable: Disable progress bar
        s: Number of sub-steps
        derivative_order: Order of Taylor expansion
    """
    extra_args = {} if extra_args is None else extra_args
    
    sampler = STORKSampler(solver_order=4, s=s, derivative_order=derivative_order)
    
    # Check if coefficients loaded successfully
    if sampler.ms is None:
        print("Warning: STORK4 coefficients not loaded, falling back to STORK2")
        return sample_stork2(model, x, sigmas, extra_args, callback, disable, s, derivative_order, debug)
    
    if debug:
        print(f"\n=== STORK4 Sampler Debug ===")
        print(f"Input shape: {x.shape}")
        print(f"Sigmas range: [{sigmas.min():.6f}, {sigmas.max():.6f}]")
        print(f"Number of steps: {len(sigmas) - 1}")
        print(f"Substeps (s): {s}")
        print(f"Derivative order: {derivative_order}")
    
    s_in = x.new_ones([x.shape[0]])
    
    for i in range(len(sigmas) - 1):
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        
        if sigma_next == 0:
            break
            
        # Get model output
        denoised = model(x, sigma * s_in, **extra_args)
        
        # Convert to ODE derivative
        d = (x - denoised) / append_dims(sigma, x.ndim)
        
        if debug and i < 3:
            print(f"\nStep {i}: sigma={sigma:.6f}, sigma_next={sigma_next:.6f}")
            print(f"  d stats: mean={d.mean():.4f}, std={d.std():.4f}")
        
        # Startup phase
        if i < sampler.derivative_order:
            sampler.model_outputs.append(d)
            dt = sigmas[i + 1] - sigmas[i]
            x = x + d * dt
        else:
            # Main STORK4 step
            dt = sigma_next - sigma  # dt should be negative (going from high to low sigma)
            
            # Calculate derivatives
            h1 = sigmas[i-1] - sigmas[i]
            h2 = sigmas[i-2] - sigmas[i-1] if i >= 2 else h1
            h3 = sigmas[i-3] - sigmas[i-2] if i >= 3 else h2
            
            outputs = sampler.model_outputs
            
            if sampler.derivative_order >= 2 and len(outputs) >= 2:
                derivative = (-outputs[-2] + 4 * outputs[-1] - 3 * d) / (2 * h1)
                second_derivative = 2 / (h1 * h2 * (h1 + h2)) * (outputs[-2] * h1 - outputs[-1] * (h1 + h2) + d * h2)
                
                if sampler.derivative_order == 3 and len(outputs) >= 3:
                    third_derivative = 6 * ((h2 - h3) * (outputs[-1] - d) + (h3 - h1) * (outputs[-2] - d) + (h1 - h2) * (outputs[-3] - d)) / (h1 * h2 * h3)
                else:
                    third_derivative = torch.zeros_like(d)
            else:
                derivative = torch.zeros_like(d)
                second_derivative = torch.zeros_like(d)
                third_derivative = torch.zeros_like(d)
            
            sampler.model_outputs.append(d)
            
            # STORK4 implementation
            mdeg, mp = sampler.mdegr(sampler.s, sampler.ms)
            mz = int(mp[0])
            mr = int(mp[1])
            
            Y_j_2 = x
            Y_j_1 = x
            Y_j = x
            
            ci1 = torch.zeros_like(d)
            ci2 = torch.zeros_like(d)
            ci3 = torch.zeros_like(d)
            
            # First part of STORK4 update
            for j in range(1, mdeg + 1):
                if j == 1:
                    temp1 = dt * sampler.recf[mr]
                    ci1 = temp1
                    ci2 = ci1
                    Y_j_2 = x
                    Y_j_1 = x + temp1 * d
                else:
                    diff = ci1
                    velocity = sampler.taylor_approximation(sampler.derivative_order, diff, d, derivative, second_derivative, third_derivative)
                    
                    temp1 = -dt * sampler.recf[mr + 2 * (j-2) + 1]
                    temp3 = -sampler.recf[mr + 2 * (j-2) + 2]
                    temp2 = 1 - temp3
                    
                    ci1 = temp1 + temp2 * ci2 + temp3 * ci3
                    Y_j = temp1 * velocity + temp2 * Y_j_1 + temp3 * Y_j_2
                
                Y_j_2 = Y_j_1
                Y_j_1 = Y_j
                ci3 = ci2
                ci2 = ci1
            
            # Four-step finishing procedure
            # Step 1
            temp1 = dt * sampler.fpa[mz, 0]
            diff = ci1
            velocity = sampler.taylor_approximation(sampler.derivative_order, diff, d, derivative, second_derivative, third_derivative)
            Y_j_1 = velocity
            Y_j_3 = Y_j + temp1 * Y_j_1
            
            # Step 2
            ci2 = ci1 + temp1
            temp1 = dt * sampler.fpa[mz, 1]
            temp2 = dt * sampler.fpa[mz, 2]
            diff = ci2
            velocity = sampler.taylor_approximation(sampler.derivative_order, diff, d, derivative, second_derivative, third_derivative)
            Y_j_2 = velocity
            Y_j_4 = Y_j + temp1 * Y_j_1 + temp2 * Y_j_2
            
            # Step 3
            ci2 = ci1 + temp1 + temp2
            temp1 = dt * sampler.fpa[mz, 3]
            temp2 = dt * sampler.fpa[mz, 4]
            temp3 = dt * sampler.fpa[mz, 5]
            diff = ci2
            velocity = sampler.taylor_approximation(sampler.derivative_order, diff, d, derivative, second_derivative, third_derivative)
            Y_j_3 = velocity
            fnt = Y_j + temp1 * Y_j_1 + temp2 * Y_j_2 + temp3 * Y_j_3
            
            # Step 4
            ci2 = ci1 + temp1 + temp2 + temp3
            temp1 = dt * sampler.fpb[mz, 0]
            temp2 = dt * sampler.fpb[mz, 1]
            temp3 = dt * sampler.fpb[mz, 2]
            temp4 = dt * sampler.fpb[mz, 3]
            diff = ci2
            velocity = sampler.taylor_approximation(sampler.derivative_order, diff, d, derivative, second_derivative, third_derivative)
            Y_j_4 = velocity
            Y_j = Y_j + temp1 * Y_j_1 + temp2 * Y_j_2 + temp3 * Y_j_3 + temp4 * Y_j_4
            
            x = Y_j
        
        # Keep only recent outputs
        if len(sampler.model_outputs) > sampler.derivative_order + 1:
            sampler.model_outputs.pop(0)
            
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigma, 'sigma_hat': sigma, 'denoised': denoised})
        
        # Debug the progression every 5 steps
        if debug and i % 5 == 0:
            print(f"\nProgress - Step {i}/{len(sigmas)-1}:")
            print(f"  Current x stats: mean={x.mean():.4f}, std={x.std():.4f}, min={x.min():.4f}, max={x.max():.4f}")
            
    if debug:
        print(f"\nFinal output:")
        print(f"  x stats: mean={x.mean():.4f}, std={x.std():.4f}, min={x.min():.4f}, max={x.max():.4f}")
            
    return x